#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版只保留一个实例脚本
避免长时间等待，快速执行
"""

import requests
import json
import time
import sys
import argparse
from typing import List, Dict, Any, Optional

class SimpleKeepOneManager:
    """简化版只保留一个实例管理器"""
    
    def __init__(self, base_url: str = "http://localhost:8080"):
        self.base_url = base_url
    
    def get_servers(self) -> List[Dict[str, Any]]:
        """获取所有服务器列表"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/servers", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    return data.get('data', [])
            return []
        except:
            return []
    
    def create_server(self, region: str = "ap-beijing", gpu_type: str = "basic") -> Optional[str]:
        """创建新服务器"""
        try:
            print(f"🚀 创建新服务器: {region} - {gpu_type}")
            response = requests.post(f"{self.base_url}/api/v1/servers", 
                                   json={'region': region, 'gpu_type': gpu_type}, 
                                   timeout=30)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    instance_id = data.get('instance_id')
                    print(f"✅ 服务器创建成功: {instance_id}")
                    return instance_id
                else:
                    print(f"❌ 服务器创建失败: {data.get('error')}")
                    return None
            else:
                print(f"❌ 创建请求失败: {response.status_code}")
                return None
        except Exception as e:
            print(f"❌ 创建异常: {e}")
            return None
    
    def destroy_server(self, instance_id: str) -> bool:
        """销毁服务器"""
        try:
            print(f"🗑️ 销毁服务器: {instance_id}")
            response = requests.delete(f"{self.base_url}/api/v1/servers/{instance_id}", timeout=15)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    print(f"✅ 服务器销毁成功: {instance_id}")
                    return True
                else:
                    print(f"❌ 服务器销毁失败: {data.get('error')}")
                    return False
            else:
                print(f"❌ 销毁请求失败: {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ 销毁异常: {e}")
            return False
    
    def select_server_to_keep(self, servers: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """选择要保留的服务器 - 只保留RUNNING状态的实例"""
        if not servers:
            return None
        
        # 只选择运行中的服务器
        running_servers = [s for s in servers if s.get('status') == 'RUNNING']
        if running_servers:
            # 选择最晚创建的运行中服务器（保留最新的）
            return max(running_servers, key=lambda x: x.get('created_time', ''))
        
        # 如果没有运行中的服务器，返回None，需要等待
        return None
    
    def has_pending_servers(self, servers: List[Dict[str, Any]]) -> bool:
        """检查是否有需要等待状态的服务器"""
        if not servers:
            return False
        
        pending_statuses = ['PENDING', 'STOPPING']
        for server in servers:
            if server.get('status') in pending_statuses:
                return True
        return False
    
    def wait_for_pending_servers(self, max_wait_time: int = 300) -> bool:
        """等待PENDING状态的服务器转变"""
        print("⏳ 检测到PENDING状态的服务器，等待状态转变...")
        start_time = time.time()
        
        while time.time() - start_time < max_wait_time:
            servers = self.get_servers()
            if not servers:
                return False
            
            if not self.has_pending_servers(servers):
                print("✅ PENDING状态服务器已转变")
                return True
            
            # 显示当前状态
            pending_count = len([s for s in servers if s.get('status') in ['PENDING', 'STOPPING']])
            running_count = len([s for s in servers if s.get('status') == 'RUNNING'])
            print(f"⏳ 等待中... PENDING: {pending_count}, RUNNING: {running_count}")
            
            time.sleep(10)
        
        print("⏰ 等待PENDING状态超时")
        return False
    
    def keep_one_instance(self, preferred_region: str = "ap-beijing", preferred_gpu_type: str = "basic") -> bool:
        """只保留一个实例 - 简化版"""
        print("🎯 开始执行：只保留一个实例（简化版）")
        print("=" * 50)
        
        # 获取当前服务器列表
        servers = self.get_servers()
        if servers is None:
            print("❌ 无法获取服务器列表")
            return False
        
        current_count = len(servers)
        print(f"📊 当前服务器数量: {current_count}")
        
        # 检查是否有PENDING状态的服务器，如果有则等待
        if self.has_pending_servers(servers):
            if not self.wait_for_pending_servers():
                print("❌ 等待PENDING状态超时")
                return False
            
            # 重新获取服务器列表
            servers = self.get_servers()
            if not servers:
                print("❌ 等待后无法获取服务器列表")
                return False
            current_count = len(servers)
            print(f"📊 等待后服务器数量: {current_count}")
        
        if current_count == 0:
            # 没有实例，创建一个
            print("🚀 没有实例，创建新实例...")
            instance_id = self.create_server(preferred_region, preferred_gpu_type)
            if instance_id:
                print("✅ 成功创建实例（启动中）")
                return True
            else:
                print("❌ 实例创建失败")
                return False
        
        # 检查RUNNING状态的实例数量
        running_servers = [s for s in servers if s.get('status') == 'RUNNING']
        running_count = len(running_servers)
        
        if running_count == 1:
            print("✅ 已经只有一个RUNNING状态的实例，无需调整")
            return True
        elif running_count == 0:
            print("⚠️ 没有RUNNING状态的实例，需要创建一个")
            instance_id = self.create_server(preferred_region, preferred_gpu_type)
            if instance_id:
                print("✅ 成功创建新实例")
                return True
            else:
                print("❌ 创建新实例失败")
                return False
        else:
            # 有多个RUNNING实例，保留一个，销毁其他的
            print(f"🗑️ 有 {running_count} 个RUNNING实例，需要保留1个，销毁 {running_count-1} 个")
            
            # 选择要保留的服务器（最晚创建的）
            server_to_keep = self.select_server_to_keep(servers)
            if not server_to_keep:
                print("❌ 无法选择要保留的服务器")
                return False
            
            keep_instance_id = server_to_keep.get('instance_id')
            print(f"✅ 选择保留服务器: {keep_instance_id}")
            
            # 销毁其他RUNNING状态的服务器
            destroyed_count = 0
            for server in servers:
                instance_id = server.get('instance_id')
                status = server.get('status')
                if instance_id != keep_instance_id and status == 'RUNNING':
                    if self.destroy_server(instance_id):
                        destroyed_count += 1
                    time.sleep(1)  # 短暂等待避免并发
            
            print(f"✅ 成功销毁 {destroyed_count} 个RUNNING实例")
            print("⏳ 等待系统稳定...")
            time.sleep(3)
            
            # 简单验证
            final_servers = self.get_servers()
            final_running_count = len([s for s in final_servers if s.get('status') == 'RUNNING']) if final_servers else 0
            print(f"📊 最终RUNNING实例数量: {final_running_count}")
            
            if final_running_count >= 1:
                print("✅ 成功保留至少一个RUNNING实例")
                return True
            else:
                print("⚠️ 没有RUNNING实例了，需要创建一个")
                instance_id = self.create_server(preferred_region, preferred_gpu_type)
                if instance_id:
                    print("✅ 成功创建新实例")
                    return True
                else:
                    print("❌ 创建新实例失败")
                    return False
    
    def show_current_status(self):
        """显示当前状态"""
        print("📊 当前实例状态")
        print("=" * 50)
        
        servers = self.get_servers()
        if servers is None:
            print("❌ 无法获取服务器列表")
            return
        
        current_count = len(servers)
        print(f"📊 当前服务器数量: {current_count}")
        
        if current_count == 0:
            print("⚠️ 没有运行中的实例")
        elif current_count == 1:
            print("✅ 正好有一个实例")
        else:
            print(f"⚠️ 有 {current_count} 个实例，建议只保留一个")
        
        if servers:
            print("\n📋 实例列表:")
            for i, server in enumerate(servers, 1):
                instance_id = server.get('instance_id', 'Unknown')
                status = server.get('status', 'Unknown')
                region = server.get('region', 'Unknown')
                gpu_type = server.get('gpu_type', 'Unknown')
                print(f"  {i:2d}. {instance_id} | {status} | {region} | {gpu_type}")

def main():
    parser = argparse.ArgumentParser(description='简化版只保留一个实例脚本')
    parser.add_argument('action', choices=['keep-one', 'status'], 
                       help='操作类型: keep-one(保留一个实例), status(查看状态)')
    parser.add_argument('--region', default='ap-beijing', 
                       help='首选地域 (默认: ap-beijing)')
    parser.add_argument('--gpu-type', default='basic', 
                       help='首选GPU类型 (默认: basic)')
    parser.add_argument('--url', default='http://localhost:8080', 
                       help='HAI优化系统URL')
    
    args = parser.parse_args()
    
    manager = SimpleKeepOneManager(args.url)
    
    if args.action == 'keep-one':
        success = manager.keep_one_instance(args.region, args.gpu_type)
        if success:
            print("\n🎉 操作完成：成功保留一个实例")
            sys.exit(0)
        else:
            print("\n❌ 操作失败：无法保留一个实例")
            sys.exit(1)
    elif args.action == 'status':
        manager.show_current_status()

if __name__ == '__main__':
    main()
